import os
import sys
import pandas as pd
import requests
import shutil
from datetime import datetime
import concurrent.futures

# 获取程序运行目录
if getattr(sys, 'frozen', False):
    # 如果是打包后的程序
    current_dir = os.path.dirname(sys.executable)
else:
    # 如果是开发环境
    current_dir = os.path.dirname(os.path.abspath(__file__))

# 将目录和文件名结合，获取文件的完整路径
file_path = os.path.join(current_dir, 'down.xlsx')

# 确保文件存在
if not os.path.exists(file_path):
    print(f"错误：找不到文件 {file_path}")
    print(f"请确保 down.xlsx 文件与程序在同一目录下")
    input("按回车键退出...")
    sys.exit(1)

# 读取文件
try:
    df = pd.read_excel(file_path)
except Exception as e:
    print(f"读取文件时出错：{e}")
    input("按回车键退出...")
    sys.exit(1)

# 初始化任务数和已完成任务数
total_tasks = len(df)
completed_tasks = 0

# 初始化失败列表
failed_downloads = []

def get_folder_structure_choice():
    while True:
        print("\n请选择文件夹生成方式：")
        print("1. 按时间生成文件夹")
        print("2. 按商家名称生成文件夹")
        choice = input("请输入选择（1 或 2）: ")
        if choice in ['1', '2']:
            return choice
        print("无效的选择，请重新输入")

def parse_date(visit_time):
    if isinstance(visit_time, datetime):
        # 如果 visit_time 是日期对象，转换为字符串后截取前10位
        return visit_time.strftime('%Y-%m-%d')
    elif isinstance(visit_time, str):
        # 如果 visit_time 是字符串，直接截取前10位
        return visit_time[:10]

# 定义下载图片的函数
def download_image(row):
    global completed_tasks

    index, data = row

    # 获取图片 url、订单号和就诊时间
    image_url = data['图片']  # 替换为你的列名
    order_number = data['订单单号']  # 替换为你的列名
    visit_time = data['就诊时间']  # 替换为你的列名
    store_name = data['商家']  # 替换为你的列名

    # 如果 url 为空、order_number 为空或 visit_time 为空，order_number 为 \N 则跳过
    if pd.isnull(image_url) or pd.isnull(order_number) or pd.isnull(visit_time) or order_number == '\\N':
        return

    # 根据用户选择决定文件夹名称
    if folder_structure_choice == '1':
        folder_name = parse_date(visit_time)
    else:
        folder_name = store_name if not pd.isnull(store_name) else "未知商家"

    folder_path = os.path.join(current_dir, folder_name)

    # 创建文件夹（如果不存在）
    if not os.path.exists(folder_path):
        os.makedirs(folder_path, exist_ok=True)

    # 下载图片
    try:
        response = requests.get(image_url, stream=True)
    except requests.exceptions.RequestException as e:
        print(f"\nError downloading image from {image_url}: {e}")
        failed_downloads.append([order_number, image_url])
        return

    # 检查状态码
    if response.status_code == 200:
        # 图片的保存路径
        path = os.path.join(folder_path, f'{order_number}.jpg')
        # 以二进制写模式打开一个新的文件
        with open(path, 'wb') as out_file:
            # 将下载的图片内容写入新文件
            shutil.copyfileobj(response.raw, out_file)
        del response
    else:
        failed_downloads.append([order_number, image_url])

    # 更新已完成任务数，并打印进度
    completed_tasks += 1
    # 计算进度百分比
    progress = completed_tasks / total_tasks
    # 计算已完成的进度条长度
    completed_length = int(50 * progress)
    # 创建进度条字符串
    progress_bar = '█' * completed_length + ' ' * (50 - completed_length)
    # 打印进度条，'\r' 使光标回到行首，end='' 防止打印新的一行
    print(f"\rProgress: [{progress_bar}] {progress * 100:.2f}%", end='')

try:
    # 在开始下载之前获取用户选择
    folder_structure_choice = get_folder_structure_choice()

    max_workers = os.cpu_count() * 2

    print(f"\n开始下载，共 {total_tasks} 张图片，最大线程数为 {max_workers} 个")

    with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
        list(executor.map(download_image, df.iterrows()))

    # 如果失败列表不为空，写入本地csv文件
    if failed_downloads:
        df_failed = pd.DataFrame(failed_downloads, columns=['order_number', 'url'])
        failed_file = os.path.join(current_dir, 'failed_downloads.csv')
        df_failed.to_csv(failed_file, index=False, encoding='utf-8-sig')
        print(f"\n有 {len(failed_downloads)} 个文件下载失败，详情已写入 failed_downloads.csv")

    print("\n下载完成")

except Exception as e:
    print(f"\n程序运行出错：{e}")
finally:
    input("\n按回车键退出...")